import os
import pandas as pd
import numpy as np
from tqdm import tqdm

def process_comm_data(base_dir):
    # 创建汇总表格的数据容器
    summary_data = []
    total_sum = 0
    processed_files = 0
    
    # 只提取base_dir的末尾部分作为标识
    base_dir_name = os.path.basename(os.path.normpath(base_dir))
    
    # 遍历base_dir下的所有cn文件夹
    for cn_dir in tqdm(os.listdir(base_dir), desc="Processing cn folders"):
        cn_path = os.path.join(base_dir, cn_dir)
        
        # 识别两种类型的cn文件夹
        is_cn_dir = cn_dir.startswith("cn[") or (cn_dir.startswith("cn") and (cn_dir[2:].isdigit() or "-" in cn_dir[2:]))
        
        if not is_cn_dir or not os.path.isdir(cn_path):
            continue
            
        # 遍历cn文件夹下的所有xnode-yproc文件夹
        for config_dir in os.listdir(cn_path):
            config_path = os.path.join(cn_path, config_dir)
            
            # 检查是否是以节点数开头的配置文件夹
            prefix_options = ("1node", "2node", "4node", "8node", "16node", "32node", 
                             "64node", "128node", "256node", "512node", "1024node")
            
            if not config_dir.startswith(prefix_options) or not os.path.isdir(config_path):
                continue
                
            # 检查并处理log-0_processed.csv文件
            csv_path = os.path.join(config_path, "log-0_processed.csv")
            
            if os.path.exists(csv_path):
                try:
                    # 读取CSV文件
                    df = pd.read_csv(csv_path)
                    
                    # 第一步：删除任何只包含"comm_type"和"comm_time(us)"值的行
                    # 条件1：这两列都有值
                    has_comm_values = (~df['comm_type'].isna() & ~df['comm_time(us)'].isna())
                    
                    # 条件2：其他所有列都为空
                    other_columns = [col for col in df.columns if col not in ['comm_type', 'comm_time(us)']]
                    other_empty = df[other_columns].isna().all(axis=1) | (df[other_columns] == '').all(axis=1)
                    
                    # 同时满足两个条件的行需要删除
                    rows_to_drop = has_comm_values & other_empty
                    
                    if rows_to_drop.any():
                        print(f"在 {csv_path} 中删除 {rows_to_drop.sum()} 行仅包含comm_type和comm_time(us)的行")
                        df = df[~rows_to_drop].reset_index(drop=True)
                    
                    # 计算当前文件中comm_type=51的comm_time总和
                    comm_type_51_mask = df['comm_type'] == 51
                    file_sum = df.loc[comm_type_51_mask, 'comm_time(us)'].sum()
                    total_sum += file_sum
                    
                    # 添加汇总数据
                    summary_data.append({
                        "base_dir": base_dir_name,
                        "cn_dir": cn_dir,
                        "node_dir": config_dir,
                        "comm_type": 51,
                        "total_comm_time": file_sum
                    })
                    
                    # 如果列存在则删除
                    if 'comm_type_51_total_comm_time' in df.columns:
                        df.drop(columns=['comm_type_51_total_comm_time'], inplace=True)
                    
                    # 添加新行
                    new_row = {col: None for col in df.columns}
                    new_row['comm_type'] = 51
                    new_row['comm_time(us)'] = file_sum
                    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
                    
                    # 保存回原文件
                    df.to_csv(csv_path, index=False)
                    processed_files += 1
                    
                except Exception as e:
                    print(f"Error processing {csv_path}: {str(e)}")
    
    # 创建汇总表格
    if summary_data:
        summary_df = pd.DataFrame(summary_data)
        
        # 保存汇总表到base_dir文件夹
        summary_path = os.path.join(base_dir, "comm_type_51_summary.csv")
        summary_df.to_csv(summary_path, index=False)
        print(f"汇总表已保存到: {summary_path}")
    else:
        print("未找到任何符合条件的CSV文件，无法创建汇总表")
    
    # 返回统计信息
    return {
        "total_comm_type_51_sum": total_sum,
        "processed_files": processed_files,
        "summary_path": summary_path if summary_data else None
    }

# 使用示例
if __name__ == "__main__":
    base_dir = r"F:/PostGraduate/Point-to-Point-DATA/1-16nodes-data"
    results = process_comm_data(base_dir)
    
    print("\n" + "="*50)
    print(f"总处理文件数: {results['processed_files']}")
    print(f"所有文件中comm_type=51的总和: {results['total_comm_type_51_sum']} (us)")
    
    if results['summary_path']:
        print(f"汇总表路径: {results['summary_path']}")
        # 打印汇总表内容
        summary_df = pd.read_csv(results['summary_path'])
        print("\n汇总表内容:")
        print(summary_df)
    
    print("="*50)